import logging
import os
import cv2
import numpy as np
from tqdm import tqdm
from staintools import StainNormalizer
import staintools

from util.log_util import get_logger

# ===== 参数设置 =====
input_patch_dir = "D:/zm/PGGN"  # 原始Patch文件夹
output_patch_dir = "D:/zm/PGGN_Pic"  # 处理后保存路径

patch_size = 256
white_threshold = 0.8  # 空白像素比例阈值

# ===== 加载颜色标准化参考图像（必须选一张代表性patch）=====
target_img_path = "D:/zm/PGGN/patch_3840_42752.png"  # 参考图像路径（来自一个染色正常的 Patch）


# true： 图像中有组织成分，图像不是空白图（比如不是全白、全黑），可以保留用于后续处理（如颜色标准化）。
def is_tissue(img_rgb, threshold=0.85):
    """
    判断是否为有组织区域：非白像素比例小于一定值就认为是空白图
    """
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    white_pixels = np.sum(img_gray > 220)  # 白色像素多
    total_pixels = img_gray.size
    white_ratio = white_pixels / total_pixels
    return white_ratio < threshold  # 非白色区域多 → 是组织


# 排除低对比度图像（如模糊、过曝、全灰等），这些图像会导致染色标准化失败
# 进一步验证图像是否适合进行染色标准化处理，正常组织图像的标准差通常 >20，模糊/低对比度图像的标准差 5~15。
def is_valid_for_staining(img_rgb):
    if np.std(img_rgb) < 10:  # 计算RGB三通道的总体标准差
        return False  # 标准差＜10 → 图像对比度过低 → 无效
    return True  # 否则有效


if __name__ == "__main__":
    # 初始化日志格式
    get_logger(logging.INFO)

    os.makedirs(output_patch_dir, exist_ok=True)
    target = cv2.imread(target_img_path)
    target = cv2.cvtColor(target, cv2.COLOR_BGR2RGB)
    normalizer = StainNormalizer(method='macenko')  # 可选：reinhard, macenko, vahadane
    normalizer.fit(target)

    # ===== 开始处理所有Patch =====
    image_num = 0
    patch_files = [f for f in os.listdir(input_patch_dir) if f.endswith('.png')]
    for fname in tqdm(patch_files):
        fpath = os.path.join(input_patch_dir, fname)
        img = cv2.imread(fpath)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # np.mean(img_rgb): 专门排除黑图，黑图虽然非白区域很多，但没有组织信息。（有组织的，肯定非白区域很多。但非白区域多一定是组织？ 不，可能是黑图）
        # is_tissue(img_rgb): 用来排除“全白图/大面积空白”，因为没有组织。
        if np.mean(img_rgb) < 5 or not is_tissue(img_rgb) or not is_valid_for_staining(img_rgb):  # 纯黑图像的均值非常低
            logging.info("图片：%s 不合规，跳过。全黑图：%s, 没有组织（白图）：%s，不合适进行标准化：%s", fname, np.mean(img_rgb) < 5, not is_tissue(img_rgb), not is_valid_for_staining(img_rgb))
            continue  # 跳过

        # ---------- 1. 空白区域检测 ----------
        gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
        white_ratio = np.sum(gray > 220) / (patch_size * patch_size)
        if white_ratio > white_threshold:
            continue  # 跳过空白图

        # ---------- 2. 颜色标准化 ----------
        try:
            img_norm = normalizer.transform(img_rgb)
            if np.isnan(img_norm).any() or np.mean(img_norm) < 5:
                logging.warning("标准化后图像异常，跳过：%s", fname)
                continue
            img_norm_bgr = cv2.cvtColor(img_norm, cv2.COLOR_RGB2BGR)
            cv2.imwrite(os.path.join(output_patch_dir, fname), img_norm_bgr)
        except Exception as e:
            logging.error("标准化失败：%s， 错误: %s", fname, str(e))
        image_num = image_num + 1
        if image_num % 500 == 0:
            logging.warning("已经处理图片: %s 张", image_num)
    logging.info("共计图片: %s, 有效图片数量: %s", len(patch_files), image_num)
